/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2024, Individual contributors, see AUTHORS file     |
   | See: https://www.mrpt.org/Authors - All rights reserved.               |
   | Released under BSD License. See: https://www.mrpt.org/License          |
   +------------------------------------------------------------------------+ */

#include "maps-precomp.h"  // Precomp header
#include "mrpt/math/ops_containers.h"
//
#include <mrpt/maps/CColouredPointsMap.h>
#include <mrpt/obs/CSensoryFrame.h>
#include <mrpt/obs/customizable_obs_viz.h>
#include <mrpt/opengl/CAxis.h>
#include <mrpt/opengl/CPlanarLaserScan.h>
#include <mrpt/opengl/CPointCloudColoured.h>
#include <mrpt/opengl/stock_objects.h>

using namespace mrpt::obs;

void VisualizationParameters::save_to_ini_file(
    mrpt::config::CConfigFileBase& cfg, const std::string& section) const
{
  auto& c = cfg;
  const std::string& s = section;

  MRPT_SAVE_CONFIG_VAR(axisTickFrequency, c, s);
  MRPT_SAVE_CONFIG_VAR(axisLimits, c, s);
  MRPT_SAVE_CONFIG_VAR(axisTickTextSize, c, s);
  MRPT_SAVE_CONFIG_VAR(colorFromRGBimage, c, s);
  MRPT_SAVE_CONFIG_VAR(colorizeByField, c, s);
  MRPT_SAVE_CONFIG_VAR(invertColorMapping, c, s);
  MRPT_SAVE_CONFIG_VAR(pointSize, c, s);
  MRPT_SAVE_CONFIG_VAR(colorMap, c, s);
  MRPT_SAVE_CONFIG_VAR(drawSensorPose, c, s);
  MRPT_SAVE_CONFIG_VAR(sensorPoseScale, c, s);
  MRPT_SAVE_CONFIG_VAR(showAxis, c, s);
  MRPT_SAVE_CONFIG_VAR(showSurfaceIn2Dscans, c, s);
  MRPT_SAVE_CONFIG_VAR(showPointsIn2Dscans, c, s);
  MRPT_SAVE_CONFIG_VAR(onlyPointsWithColor, c, s);

  MRPT_SAVE_CONFIG_VAR(surface2DscansColor.R, c, s);
  MRPT_SAVE_CONFIG_VAR(surface2DscansColor.G, c, s);
  MRPT_SAVE_CONFIG_VAR(surface2DscansColor.B, c, s);
  MRPT_SAVE_CONFIG_VAR(surface2DscansColor.A, c, s);

  MRPT_SAVE_CONFIG_VAR(points2DscansColor.R, c, s);
  MRPT_SAVE_CONFIG_VAR(points2DscansColor.G, c, s);
  MRPT_SAVE_CONFIG_VAR(points2DscansColor.B, c, s);
  MRPT_SAVE_CONFIG_VAR(points2DscansColor.A, c, s);
}

void VisualizationParameters::load_from_ini_file(
    const mrpt::config::CConfigFileBase& cfg, const std::string& section)
{
  const auto& c = cfg;
  const std::string& s = section;

  MRPT_LOAD_CONFIG_VAR_CS(axisTickFrequency, double);
  MRPT_LOAD_CONFIG_VAR_CS(axisLimits, double);
  MRPT_LOAD_CONFIG_VAR_CS(axisTickTextSize, double);
  MRPT_LOAD_CONFIG_VAR_CS(colorFromRGBimage, bool);
  MRPT_LOAD_CONFIG_VAR_CS(colorizeByField, string);
  MRPT_LOAD_CONFIG_VAR_CS(invertColorMapping, bool);
  MRPT_LOAD_CONFIG_VAR_CS(pointSize, double);
  MRPT_LOAD_CONFIG_VAR_CS(drawSensorPose, bool);
  MRPT_LOAD_CONFIG_VAR_CS(sensorPoseScale, double);
  colorMap = c.read_enum(s, "colorMap", colorMap);
  MRPT_LOAD_CONFIG_VAR_CS(showAxis, bool);
  MRPT_LOAD_CONFIG_VAR_CS(showSurfaceIn2Dscans, bool);
  MRPT_LOAD_CONFIG_VAR_CS(showPointsIn2Dscans, bool);
  MRPT_LOAD_CONFIG_VAR_CS(onlyPointsWithColor, bool);

  MRPT_LOAD_CONFIG_VAR_CS(surface2DscansColor.R, int);
  MRPT_LOAD_CONFIG_VAR_CS(surface2DscansColor.G, int);
  MRPT_LOAD_CONFIG_VAR_CS(surface2DscansColor.B, int);
  MRPT_LOAD_CONFIG_VAR_CS(surface2DscansColor.A, int);

  MRPT_LOAD_CONFIG_VAR_CS(points2DscansColor.R, int);
  MRPT_LOAD_CONFIG_VAR_CS(points2DscansColor.G, int);
  MRPT_LOAD_CONFIG_VAR_CS(points2DscansColor.B, int);
  MRPT_LOAD_CONFIG_VAR_CS(points2DscansColor.A, int);
}

// Bounding box memory so we have consistent coloring across different sensors:
namespace
{
thread_local std::map<std::string, std::pair<float, float>> bbMemory;
constexpr double bbMemoryFading = 0.99;
}  // namespace

void mrpt::obs::recolorize3Dpc(
    const mrpt::opengl::CPointCloudColoured::Ptr& pnts,
    const mrpt::maps::CPointsMap* originalPts,
    const VisualizationParameters& p)
{
  if (!originalPts)
  {
    return;
  }

  std::pair<float, float> dataLimits{0, 0};
  std::size_t dataPoints = 0;

  const mrpt::aligned_std_vector<float>* fieldData_f = nullptr;
  const mrpt::aligned_std_vector<uint16_t>* fieldData_u = nullptr;

  if (fieldData_f = originalPts->getPointsBufferRef_float_field(p.colorizeByField); fieldData_f)
  {
    dataPoints = fieldData_f->size();
    if (dataPoints)
    {
      mrpt::math::minimum_maximum(*fieldData_f, dataLimits.first, dataLimits.second);
    }
  }
  else if (fieldData_u = originalPts->getPointsBufferRef_uint_field(p.colorizeByField); fieldData_u)
  {
    dataPoints = fieldData_u->size();

    if (dataPoints)
    {
      uint16_t uMin = 0, uMax = 0;
      mrpt::math::minimum_maximum(*fieldData_u, uMin, uMax);
      dataLimits.first = static_cast<float>(uMin);
      dataLimits.second = static_cast<float>(uMax);
    }
  }

  if (!fieldData_f && !fieldData_u)
  {
    return;
  }

  if (dataPoints != pnts->size())
  {
    std::cerr << "[mrpt::obs::recolorize3Dpc] Warning: Skipping colorization due to mismatch in "
                 "vector lengths.\n";
    return;
  }

  // Slowly update bb memory:
  if (!bbMemory.count(p.colorizeByField))
  {
    // first time:
    bbMemory[p.colorizeByField] = dataLimits;
  }
  else
  {
    auto& bb = bbMemory[p.colorizeByField];
    bb.first = bb.first * bbMemoryFading + dataLimits.first * (1.0 - bbMemoryFading);
    bb.second = bb.second * bbMemoryFading + dataLimits.second * (1.0 - bbMemoryFading);
  }
  const auto& bb = bbMemory[p.colorizeByField];

  const float colMin = p.invertColorMapping ? bb.second : bb.first;
  const float colMax = p.invertColorMapping ? bb.first : bb.second;

  const float coord_range = colMax - colMin;
  const float coord_range_1 = coord_range != 0.0f ? 1.0f / coord_range : 1.0f;
  for (size_t i = 0; i < pnts->size(); i++)
  {
    float coord = fieldData_f ? fieldData_f->at(i) : static_cast<float>(fieldData_u->at(i));
    const float col_idx = std::max(0.0f, std::min(1.0f, (coord - colMin) * coord_range_1));
    float r, g, b;
    mrpt::img::colormap(p.colorMap, col_idx, r, g, b);
    pnts->setPointColor_fast(i, r, g, b);
  }
}

namespace
{
void add_common_to_viz(
    const CObservation& obs, const VisualizationParameters& p, mrpt::opengl::CSetOfObjects& out)
{
  if (p.showAxis)
  {
    const float L = p.axisLimits;
    auto gl_axis = mrpt::opengl::CAxis::Create(-L, -L, -L, L, L, L, p.axisTickFrequency, 2, true);
    gl_axis->setTextScale(p.axisTickTextSize);
    gl_axis->setColor_u8(0xa0, 0xa0, 0xa0, 0x80);
    out.insert(gl_axis);

    // Show axis labels such that they can be read from the IIIrd-IVth
    // quadrant:
    const float yawIncrs[3] = {180.f, -90.f, 180.f};
    for (int axis = 0; axis < 3; axis++)
    {
      float yaw_deg, pitch_deg, roll_deg;
      gl_axis->getTextLabelOrientation(axis, yaw_deg, pitch_deg, roll_deg);
      yaw_deg += yawIncrs[axis];
      gl_axis->setTextLabelOrientation(axis, yaw_deg, pitch_deg, roll_deg);
    }
  }

  if (p.drawSensorPose)
  {
    const auto glCorner = mrpt::opengl::stock_objects::CornerXYZSimple(p.sensorPoseScale);
    glCorner->setPose(obs.sensorPose());
    out.insert(glCorner);
  }
}
}  // namespace

void mrpt::obs::obs3Dscan_to_viz(
    const CObservation3DRangeScan::Ptr& obs,
    const VisualizationParameters& p,
    mrpt::opengl::CSetOfObjects& out)
{
  out.clear();

  // Generate/load 3D points
  // ----------------------
  mrpt::maps::CPointsMap::Ptr pointMap;
  mrpt::maps::CColouredPointsMap::Ptr pointMapCol;
  mrpt::obs::T3DPointsProjectionParams pp;
  pp.takeIntoAccountSensorPoseOnRobot = true;
  pp.onlyPointsWithIntensityColor = p.onlyPointsWithColor;

  // Color from intensity image?
  if (p.colorFromRGBimage && obs->hasRangeImage && obs->hasIntensityImage)
  {
    pointMapCol = mrpt::maps::CColouredPointsMap::Create();
    pointMapCol->colorScheme.scheme = mrpt::maps::CColouredPointsMap::cmFromIntensityImage;

    obs->unprojectInto(*pointMapCol, pp);
    pointMap = pointMapCol;
  }
  else
  {
    // Empty point set, or load from XYZ in observation:
    pointMap = mrpt::maps::CSimplePointsMap::Create();
    if (obs->hasPoints3D)
    {
      for (size_t i = 0; i < obs->points3D_x.size(); i++)
        pointMap->insertPoint(obs->points3D_x[i], obs->points3D_y[i], obs->points3D_z[i]);
    }
    else if (obs->hasRangeImage)
    {
      obs->unprojectInto(*pointMap, pp);
    }
  }

  add_common_to_viz(*obs, p, out);

  auto gl_pnts = mrpt::opengl::CPointCloudColoured::Create();
  // Load as RGB or grayscale points:
  if (pointMapCol)
    gl_pnts->loadFromPointsMap(pointMapCol.get());
  else
  {
    gl_pnts->loadFromPointsMap(pointMap.get());
    recolorize3Dpc(gl_pnts, pointMap.get(), p);
  }

  // No need to further transform 3D points
  gl_pnts->setPose(mrpt::poses::CPose3D());
  gl_pnts->setPointSize(p.pointSize);

  out.insert(gl_pnts);
}

void mrpt::obs::obsVelodyne_to_viz(
    const CObservationVelodyneScan::Ptr& obs,
    const VisualizationParameters& p,
    mrpt::opengl::CSetOfObjects& out)
{
  out.clear();

  add_common_to_viz(*obs, p, out);

  auto pnts = mrpt::opengl::CPointCloudColoured::Create();
  out.insert(pnts);

  mrpt::maps::CColouredPointsMap pntsMap;
  pntsMap.loadFromVelodyneScan(*obs);
  pnts->loadFromPointsMap(&pntsMap);
  pnts->setPointSize(p.pointSize);

  if (!p.colorFromRGBimage)
  {
    recolorize3Dpc(pnts, &pntsMap, p);
  }
}

void mrpt::obs::obsPointCloud_to_viz(
    const mrpt::obs::CObservationPointCloud::Ptr& obs,
    const VisualizationParameters& p,
    mrpt::opengl::CSetOfObjects& out)
{
  out.clear();

  add_common_to_viz(*obs, p, out);

  auto pnts = mrpt::opengl::CPointCloudColoured::Create();
  out.insert(pnts);

  if (obs->pointcloud) pnts->loadFromPointsMap(obs->pointcloud.get());
  pnts->setPose(obs->sensorPose);

  pnts->setPointSize(p.pointSize);

  if (!p.colorFromRGBimage)
  {
    recolorize3Dpc(pnts, obs->pointcloud.get(), p);
  }
}

void mrpt::obs::obsRotatingScan_to_viz(
    const mrpt::obs::CObservationRotatingScan::Ptr& obs,
    const VisualizationParameters& p,
    mrpt::opengl::CSetOfObjects& out)
{
  out.clear();

  add_common_to_viz(*obs, p, out);

  auto pnts = mrpt::opengl::CPointCloudColoured::Create();
  out.insert(pnts);

  if (!obs->organizedPoints.empty())
  {
    for (const auto& pt : obs->organizedPoints) pnts->insertPoint({pt.x, pt.y, pt.z, 0, 0, 0, 0});
  }
  pnts->setPose(obs->sensorPose);

  pnts->setPointSize(p.pointSize);

  if (!p.colorFromRGBimage) recolorize3Dpc(pnts, nullptr, p);
}

void mrpt::obs::obs2Dscan_to_viz(
    const CObservation2DRangeScan::Ptr& obs,
    const VisualizationParameters& p,
    mrpt::opengl::CSetOfObjects& out)
{
  out.clear();

  add_common_to_viz(*obs, p, out);

  auto pnts = mrpt::opengl::CPlanarLaserScan::Create();
  out.insert(pnts);

  pnts->setScan(*obs);
  pnts->setPointSize(p.pointSize);
  pnts->enableSurface(p.showSurfaceIn2Dscans);
  pnts->enablePoints(p.showPointsIn2Dscans);

  pnts->setSurfaceColor(
      mrpt::u8tof(p.surface2DscansColor.R), mrpt::u8tof(p.surface2DscansColor.G),
      mrpt::u8tof(p.surface2DscansColor.B), mrpt::u8tof(p.surface2DscansColor.A));
  pnts->setPointsColor(
      mrpt::u8tof(p.points2DscansColor.R), mrpt::u8tof(p.points2DscansColor.G),
      mrpt::u8tof(p.points2DscansColor.B), mrpt::u8tof(p.points2DscansColor.A));
}

bool mrpt::obs::obs_to_viz(
    const mrpt::obs::CObservation::Ptr& obs,
    const VisualizationParameters& p,
    mrpt::opengl::CSetOfObjects& out)
{
  using namespace mrpt::obs;

  if (const auto o3D = std::dynamic_pointer_cast<CObservation3DRangeScan>(obs); o3D)
  {
    obs3Dscan_to_viz(o3D, p, out);
    return true;
  }
  else if (const auto oVel = std::dynamic_pointer_cast<CObservationVelodyneScan>(obs); oVel)
  {
    obsVelodyne_to_viz(oVel, p, out);
    return true;
  }
  else if (const auto oPC = std::dynamic_pointer_cast<CObservationPointCloud>(obs); oPC)
  {
    obsPointCloud_to_viz(oPC, p, out);
    return true;
  }
  else if (const auto oRS = std::dynamic_pointer_cast<CObservationRotatingScan>(obs); oRS)
  {
    obsRotatingScan_to_viz(oRS, p, out);
    return true;
  }
  else if (const auto o2D = std::dynamic_pointer_cast<CObservation2DRangeScan>(obs); o2D)
  {
    obs2Dscan_to_viz(o2D, p, out);
    return true;
  }

  // No conversion found:
  out.clear();
  return false;
}

bool mrpt::obs::obs_to_viz(
    const mrpt::obs::CSensoryFrame& sf,
    const VisualizationParameters& p,
    mrpt::opengl::CSetOfObjects& out)
{
  out.clear();

  // Make a copy, to avoid potentially duplicated axis:
  VisualizationParameters vp = p;

  for (const auto& obs : sf)
  {
    if (!obs) continue;
    auto glObj = mrpt::opengl::CSetOfObjects::Create();

    bool ok = obs_to_viz(obs, vp, *glObj);
    if (!ok) continue;

    out.insert(glObj);

    vp.showAxis = false;  // to avoid duplications
  }

  return !out.empty();
}
